import torch
import pdb
class RMSE_Loss(torch.nn.Module):
    """Root Mean Squared Error Loss with Dynamic Scaling"""
    def __init__(self, args, dataset):
        super().__init__()
        weights = torch.tensor(args.class_weights).to(args.device)
        self.weights = self.dyn_scale(args.task, dataset, weights)

    def dyn_scale(self, task, dataset, weights):
        def scale(labels):
            return weights
        return scale

    def forward(self, predictions, targets):
        ## pdb.set_trace()
        loss = 0
        alpha = 1
        num = 0
        for i in range(len(predictions)):
            prediction = predictions[i]
            target = targets[i].view(-1, 1)
            loss += (alpha * (prediction - target) ** 2).sum()
            num += target.size(0)
        #alpha = self.weights(targets)[targets].view(-1, 1)
        loss = loss / num
        return loss
